"""
图像处理工具模块
提供图像预处理、后处理和格式转换等功能
"""

import cv2
import numpy as np
from PIL import Image
from typing import Tuple, Optional, Union
from src.config import config
from src.utils.logger import app_logger


class ImageProcessor:
    """
    图像处理器类
    
    用途：
    - 图像预处理（缩放、裁剪、增强等）
    - 图像后处理（滤波、去噪等）
    - 图像格式转换
    
    使用场景：
    - 车牌检测前的图像预处理
    - 车牌识别前的图像增强
    - 结果展示的图像格式转换
    """
    
    def __init__(self):
        """
        初始化图像处理器
        """
        self.max_size = config.get('image_processing.max_image_size')
        self.min_plate_area = config.get('image_processing.min_plate_area')
        self.blur_kernel = config.get('image_processing.gaussian_blur_kernel')
    
    def load_image(self, image_path: str) -> Optional[np.ndarray]:
        """
        加载图像文件
        
        参数：
            image_path (str): 图像文件路径
            
        返回值：
            np.ndarray: 图像数组，BGR格式，失败时返回None
        """
        try:
            image = cv2.imread(image_path)
            if image is None:
                app_logger.error(f"无法加载图像文件: {image_path}")
                return None
            
            app_logger.info(f"成功加载图像: {image_path}, 尺寸: {image.shape}")
            return image
        except Exception as e:
            app_logger.error(f"加载图像文件失败: {e}")
            return None
    
    def resize_image(self, image: np.ndarray, max_size: int = None) -> np.ndarray:
        """
        调整图像尺寸，保持宽高比
        
        参数：
            image (np.ndarray): 输入图像
            max_size (int): 最大尺寸，默认使用配置值
            
        返回值：
            np.ndarray: 调整后的图像
        """
        if max_size is None:
            max_size = self.max_size
        
        height, width = image.shape[:2]
        
        # 如果图像尺寸已经小于最大尺寸，直接返回
        if max(height, width) <= max_size:
            return image
        
        # 计算缩放比例
        if height > width:
            new_height = max_size
            new_width = int(width * max_size / height)
        else:
            new_width = max_size
            new_height = int(height * max_size / width)
        
        resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
        app_logger.debug(f"图像尺寸调整: {width}x{height} -> {new_width}x{new_height}")
        
        return resized
    
    def enhance_image(self, image: np.ndarray) -> np.ndarray:
        """
        增强图像质量
        
        参数：
            image (np.ndarray): 输入图像
            
        返回值：
            np.ndarray: 增强后的图像
        """
        # 转换为LAB色彩空间
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # 应用CLAHE（对比度限制自适应直方图均衡，参数从配置文件读取）
        clip_limit = config.get('image_processing.clahe_clip_limit')
        tile_grid_size = tuple(config.get('image_processing.clahe_tile_grid_size'))
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        l = clahe.apply(l)
        
        # 合并通道并转换回BGR
        enhanced = cv2.merge([l, a, b])
        enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
        
        # 轻微锐化
        kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
        enhanced = cv2.filter2D(enhanced, -1, kernel)
        
        return enhanced
    
    def preprocess_for_detection(self, image: np.ndarray) -> np.ndarray:
        """
        为车牌检测预处理图像
        
        参数：
            image (np.ndarray): 输入图像
            
        返回值：
            np.ndarray: 预处理后的图像
        """
        # 调整尺寸
        processed = self.resize_image(image)
        
        # 增强图像
        processed = self.enhance_image(processed)
        
        # 轻微去噪
        processed = cv2.bilateralFilter(processed, 9, 75, 75)
        
        return processed
    
    def preprocess_for_ocr(self, plate_image: np.ndarray) -> np.ndarray:
        """
        为OCR识别预处理车牌图像
        
        参数：
            plate_image (np.ndarray): 车牌图像
            
        返回值：
            np.ndarray: 预处理后的车牌图像
        """
        # 对于车牌识别，简化预处理以保持字符特征
        
        # 调整尺寸，OCR通常在较大尺寸下效果更好
        height, width = plate_image.shape[:2]
        min_height = config.get('recognition.ocr_preprocessing.min_height_for_resize', 64)
        
        if height < min_height:
            scale = min_height / height
            new_width = int(width * scale)
            resized = cv2.resize(plate_image, (new_width, min_height), interpolation=cv2.INTER_CUBIC)
        else:
            # 如果已经足够大，适度放大以提高识别率
            scale = 2.0
            new_width = int(width * scale)
            new_height = int(height * scale)
            resized = cv2.resize(plate_image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
        
        # 轻微的对比度增强
        if len(resized.shape) == 3:
            # 彩色图像，转换为LAB空间增强
            lab = cv2.cvtColor(resized, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            l = clahe.apply(l)
            enhanced = cv2.merge([l, a, b])
            enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
        else:
            # 灰度图像
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            enhanced = clahe.apply(resized)
        
        # 轻微去噪，保持字符边缘
        denoised = cv2.bilateralFilter(enhanced, 9, 75, 75)
        
        return denoised
    
    def crop_plate(self, image: np.ndarray, bbox: Tuple[int, int, int, int]) -> Optional[np.ndarray]:
        """
        从图像中裁剪车牌区域
        
        参数：
            image (np.ndarray): 原始图像
            bbox (Tuple[int, int, int, int]): 边界框坐标 (x1, y1, x2, y2)
            
        返回值：
            np.ndarray: 裁剪的车牌图像，失败时返回None
        """
        try:
            x1, y1, x2, y2 = bbox
            
            # 边界检查
            height, width = image.shape[:2]
            x1 = max(0, min(x1, width - 1))
            y1 = max(0, min(y1, height - 1))
            x2 = max(x1 + 1, min(x2, width))
            y2 = max(y1 + 1, min(y2, height))
            
            # 检查车牌区域大小
            plate_area = (x2 - x1) * (y2 - y1)
            if plate_area < self.min_plate_area:
                app_logger.warning(f"车牌区域太小: {plate_area} < {self.min_plate_area}")
                return None
            
            # 裁剪车牌
            plate = image[y1:y2, x1:x2]
            
            return plate
        except Exception as e:
            app_logger.error(f"裁剪车牌失败: {e}")
            return None
    
    def draw_detection_results(self, image: np.ndarray, detections: list) -> np.ndarray:
        """
        在图像上绘制检测结果
        
        参数：
            image (np.ndarray): 原始图像
            detections (list): 检测结果列表
            
        返回值：
            np.ndarray: 绘制了检测结果的图像
        """
        result_image = image.copy()
        
        for detection in detections:
            bbox = detection.get('bbox')
            confidence = detection.get('confidence', 0.0)
            # 优先使用processed_text，如果没有则使用text
            text = detection.get('processed_text', detection.get('text', ''))
            
            if bbox is not None:
                x1, y1, x2, y2 = bbox
                
                # 绘制边界框
                cv2.rectangle(result_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # 绘制车牌号和置信度
                if text:
                    label = f"{text} ({confidence:.2f})"
                else:
                    label = f"{confidence:.2f}"
                
                # 使用PIL绘制中文文本，增大字体
                result_image = self._draw_chinese_text(result_image, label, (x1, y1 - 40), font_size=32)
        
        return result_image
    
    def _draw_chinese_text(self, image: np.ndarray, text: str, position: tuple, 
                          font_size: int = 20, color: tuple = (0, 0, 0), 
                          bg_color: tuple = (0, 255, 0)) -> np.ndarray:
        """
        在图像上绘制中文文本
        
        参数：
            image (np.ndarray): 输入图像
            text (str): 要绘制的文本
            position (tuple): 文本位置 (x, y)
            font_size (int): 字体大小
            color (tuple): 文本颜色 (R, G, B)
            bg_color (tuple): 背景颜色 (R, G, B)
            
        返回值：
            np.ndarray: 绘制了文本的图像
        """
        try:
            from PIL import Image, ImageDraw, ImageFont
            import platform
            import os
            
            # 转换为PIL图像
            pil_image = self.cv2_to_pil(image)
            draw = ImageDraw.Draw(pil_image)
            
            # 尝试加载中文字体
            font = None
            font_paths = []
            
            if platform.system() == "Windows":
                # Windows系统字体路径，优先使用宋体
                font_paths = [
                    "C:/Windows/Fonts/simsun.ttc",    # 宋体
                    "C:/Windows/Fonts/simhei.ttf",    # 黑体
                    "C:/Windows/Fonts/msyh.ttc",      # 微软雅黑
                    "C:/Windows/Fonts/simkai.ttf",    # 楷体
                    "C:/Windows/Fonts/arial.ttf",     # Arial
                ]
            elif platform.system() == "Darwin":  # macOS
                font_paths = [
                    "/System/Library/Fonts/PingFang.ttc",
                    "/System/Library/Fonts/Helvetica.ttc",
                ]
            else:  # Linux
                font_paths = [
                    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
                    "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
                ]
            
            # 尝试加载字体
            for font_path in font_paths:
                try:
                    if os.path.exists(font_path):
                        font = ImageFont.truetype(font_path, font_size)
                        app_logger.debug(f"成功加载字体: {font_path}")
                        break
                except Exception as e:
                    app_logger.debug(f"加载字体失败 {font_path}: {e}")
                    continue
            
            # 如果所有字体都失败，使用默认字体
            if font is None:
                try:
                    font = ImageFont.load_default()
                    app_logger.warning("使用默认字体")
                except:
                    # 如果连默认字体都失败，回退到OpenCV
                    return self._draw_opencv_text(image, text, position, font_size, color, bg_color)
            
            # 获取文本边界框
            try:
                bbox = draw.textbbox((0, 0), text, font=font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
            except:
                # 如果textbbox失败，使用textsize
                try:
                    text_width, text_height = draw.textsize(text, font=font)
                except:
                    # 如果都失败，估算大小
                    text_width = len(text) * font_size
                    text_height = font_size
            
            # 确保位置在图像范围内
            x, y = position
            x = max(5, min(x, pil_image.width - text_width - 10))
            y = max(5, min(y, pil_image.height - text_height - 10))
            
            # 绘制背景矩形
            bg_rect = [x - 8, y - 8, x + text_width + 8, y + text_height + 8]
            draw.rectangle(bg_rect, fill=bg_color)
            
            # 绘制文本
            draw.text((x, y), text, font=font, fill=color)
            
            # 转换回OpenCV格式
            return self.pil_to_cv2(pil_image)
            
        except Exception as e:
            app_logger.warning(f"PIL绘制中文文本失败: {e}")
            # 如果失败，回退到OpenCV绘制
            return self._draw_opencv_text(image, text, position, font_size, color, bg_color)
    
    def _draw_opencv_text(self, image: np.ndarray, text: str, position: tuple, 
                         font_size: int, color: tuple, bg_color: tuple) -> np.ndarray:
        """
        使用OpenCV绘制文本（备用方案）
        
        参数：
            image (np.ndarray): 输入图像
            text (str): 要绘制的文本
            position (tuple): 文本位置 (x, y)
            font_size (int): 字体大小
            color (tuple): 文本颜色 (R, G, B)
            bg_color (tuple): 背景颜色 (R, G, B)
            
        返回值：
            np.ndarray: 绘制了文本的图像
        """
        try:
            # 将非ASCII字符替换为可显示字符
            safe_text = ''
            for char in text:
                if ord(char) < 128:
                    safe_text += char
                else:
                    # 中文字符保持原样，让OpenCV尝试显示
                    safe_text += char
            
            # 计算字体比例
            font_scale = max(1.0, font_size / 20.0)
            thickness = max(2, int(font_size / 10))
            
            # 计算文本位置
            (text_width, text_height), baseline = cv2.getTextSize(
                safe_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness
            )
            
            x, y = position
            # 确保位置在图像范围内
            x = max(5, min(x, image.shape[1] - text_width - 10))
            y = max(text_height + 5, min(y, image.shape[0] - 10))
            
            # 绘制文本背景
            cv2.rectangle(
                image, 
                (x - 8, y - text_height - 8), 
                (x + text_width + 8, y + baseline + 8), 
                bg_color, 
                -1
            )
            
            # 绘制文本
            cv2.putText(
                image, 
                safe_text, 
                (x, y), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                font_scale, 
                color, 
                thickness
            )
            
            return image
            
        except Exception as e:
            app_logger.error(f"OpenCV绘制文本失败: {e}")
            return image
    
    def pil_to_cv2(self, pil_image: Image.Image) -> np.ndarray:
        """
        将PIL图像转换为OpenCV格式
        
        参数：
            pil_image (Image.Image): PIL图像
            
        返回值：
            np.ndarray: OpenCV格式图像
        """
        return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    
    def cv2_to_pil(self, cv2_image: np.ndarray) -> Image.Image:
        """
        将OpenCV图像转换为PIL格式
        
        参数：
            cv2_image (np.ndarray): OpenCV格式图像
            
        返回值：
            Image.Image: PIL图像
        """
        return Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))


# 全局图像处理器实例
image_processor = ImageProcessor() 